'''
Created on 16 fevr. 2011
@author: paraita

partie 4 du tp 2&3
'''

import matplotlib.pyplot as plt 
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import numpy as np
import math
import Image
from PIL import ImageMath

if __name__ == '__main__':
    SZ=256
    im = Image.new("I",(SZ,SZ))
    print im
    print "debug de l'image:",im.format, im.size, im.mode
    pixels= im.load()
    Ty= 256
    
    print "SZ/Ty:",SZ,"/",Ty,"=",SZ/Ty
    
    '''
        255 BLANC
        0 NOIR
    '''
    TOTO= (SZ/(Ty/2))
    for i in range(Ty/2):
        for j in range(SZ/Ty):
            for k in range(SZ):
                #print "plop",i,j,k
                pixels[i*TOTO+j, k] = 255
                pixels[i*TOTO+j+(SZ/Ty), k] = 0
    
    plt.figure(1)
    plt.clf()
    plt.imshow(im, cmap=plt.cm.gray_r, interpolation="Nearest",origin="lower")
    #plt.show()
    
    # fft
    F1 = np.fft.fft2(im)
    F2 = np.fft.fftshift(F1)
    
    # Calculate a 2D magnitude spectrum
    psd2D =np.abs(F2)
    fig = plt.figure(2)
    plt.clf()
    X = range(-SZ/2,SZ/2)
    Y = range(-SZ/2,SZ/2)
    X, Y = np.meshgrid(X, Y)
    ax = Axes3D(fig)
    print psd2D
    ax.plot_surface(X, Y, psd2D, rstride=1, cstride=1, cmap=cm.jet)
    
    # Inverse fft
    imm = np.fft.ifft2(F1)
    plt.figure(5)
    plt.clf()
    
    #print imm
    plt.imshow(imm.real, cmap=cm.gray)
    plt.show()
